#ifdef __KERNEL__
#include <linux/errno.h>
#include <linux/string.h>
#else
#include <errno.h>
#include <string.h>
#endif
#include "ecc.h"
#include "random.h"

/**
 * @brief SM2椭圆曲线参数
 *
 */
static const struct ecc_sm2_parms sm2_ecp = {
    .desc    = "sm2p256v1",
    .nbits   = 256,
    .fips    = 0,
    .model   = MPI_EC_WEIERSTRASS,
    .dialect = ECC_DIALECT_STANDARD,
    .p       = "0xfffffffeffffffffffffffffffffffffffffffff0000000"
         "0ffffffffffffffff",
    .a = "0xfffffffeffffffffffffffffffffffffffffffff0000000"
         "0fffffffffffffffc",
    .b = "0x28e9fa9e9d9f5e344d5a9e4bcf6509a7f39789f515ab8f9"
         "2ddbcbd414d940e93",
    .n = "0xfffffffeffffffffffffffffffffffff7203df6b21c6052"
         "b53bbf40939d54123",
    .g_x =
        "0x32c4ae2c1f1981195f9904466a39c9948fe30bbff2660be1"
        "715a4589334c74c7",
    .g_y =
        "0xbc3736a2f4f6779c59bdcee36b692153d0a9877cc62a4740"
        "02df32e52139f0a0",
    .h = 1};

struct mpi_ec_ctx sm2_ctx;

/**
 * @brief 初始化参数
 *
 * @param ec
 * @return int 成功返回0，或错误码
 */
int sm2_ec_ctx_init()
{
    const struct ecc_sm2_parms* ecp = &sm2_ecp;
    MPI                         p   = NULL;
    MPI                         a   = NULL;
    MPI                         b   = NULL;
    MPI                         x   = NULL;
    MPI                         y   = NULL;
    int                         rc  = -EINVAL;

    p = mpi_scanval(ecp->p);
    a = mpi_scanval(ecp->a);
    b = mpi_scanval(ecp->b);
    x = mpi_scanval(ecp->g_x);
    y = mpi_scanval(ecp->g_y);
    if (!p || !a || !b || !x || !y) goto free;

    rc        = -ENOMEM;
    sm2_ctx.d = mpi_new(0);
    if (!sm2_ctx.d) goto free;

    sm2_ctx.Q = mpi_point_new(0);
    if (!sm2_ctx.Q) {
        mpi_free(sm2_ctx.d);
        goto free;
    }

    sm2_ctx.t.two_inv_p = mpi_new(0);
    if (!sm2_ctx.t.two_inv_p) {
        mpi_point_release(sm2_ctx.Q);
        mpi_free(sm2_ctx.d);
        goto free;
    }

    /* mpi_ec_setup_elliptic_curve */
    sm2_ctx.G = mpi_point_new(0);
    if (!sm2_ctx.G) {
        mpi_free(sm2_ctx.d);
        mpi_free(sm2_ctx.t.two_inv_p);
        mpi_point_release(sm2_ctx.Q);
        goto free;
    }

    mpi_set(sm2_ctx.G->x, x);
    mpi_set(sm2_ctx.G->y, y);
    mpi_set_ui(sm2_ctx.G->z, 1);

    rc        = -EINVAL;
    sm2_ctx.n = mpi_scanval(ecp->n);
    if (!sm2_ctx.n) {
        mpi_free(sm2_ctx.d);
        mpi_free(sm2_ctx.t.two_inv_p);
        mpi_point_release(sm2_ctx.Q);
        mpi_point_release(sm2_ctx.G);
        goto free;
    }

    sm2_ctx.h    = ecp->h;
    sm2_ctx.name = ecp->desc;
    mpi_ec_init(&sm2_ctx, ecp->model, ecp->dialect, 0, p, a, b);

    rc = 0;
free:
    mpi_free(x);
    mpi_free(y);
    mpi_free(p);
    mpi_free(a);
    mpi_free(b);
    return rc;
}
#ifndef __KERNEL__
__attribute((constructor)) void pre_ecc_init() { (void)sm2_ec_ctx_init(); }
#endif

void sm2_ec_ctx_deinit(void)
{
    mpi_ec_deinit(&sm2_ctx);
    memset(&sm2_ctx, 0, sizeof(sm2_ctx));
}
#ifndef __KERNEL__
__attribute((destructor)) void post_ecc_deinit() { (void)sm2_ec_ctx_deinit(); }
#endif

/**
 * @brief 分配密钥内存
 *
 * @return EC_KEY* 成功返回对应地址，或NULL
 */
EC_KEY* ec_key_new()
{
    EC_KEY* key = mpimem_malloc(sizeof(EC_KEY));
    if (key) {
        key->pub  = mpi_point_new(0);
        key->priv = mpi_new(0);
    }
    return key;
}

/**
 * @brief 释放密钥内存
 *
 * @param key
 */
void ec_key_free(EC_KEY* key)
{
    if (key) {
        mpi_free(key->priv);
        mpi_point_release(key->pub);
        mpimem_free(key);
    }
}

/**
 * @brief 生成私钥
 *
 * @param key 密钥
 * @return int 成功返回0，或错误码
 */
static int ec_gen_privkey(EC_KEY* key)
{
    unsigned char random[SM2_PRIVKEY_LEN];
    MPI           d  = NULL;
    MPI           n  = NULL;
    int           rc = 0;

    n = mpi_scanval(sm2_ecp.n);
    if (!n) return -ENOMEM;

    rc = mpi_sub_ui(n, n, 1);
    if (rc) goto free_n;

    do {
        rc = random_generator(random, SM2_PRIVKEY_LEN);
        if (rc) break;

        d = mpi_read_raw_data(random, SM2_PRIVKEY_LEN);
        if (!d) break;

        /* d∈[1,n-2]; */
        if (mpi_cmp_ui(d, 0) != 0 && mpi_cmp(d, n) == -1) {
            mpi_set(key->priv, d);
            break;
        }
        mpi_free(d);
    } while (1);

    mpi_free(d);
free_n:
    mpi_free(n);
    return rc;
}

/**
 * @brief 生成公私钥对
 *
 * @param key 密钥
 * @return int 成功返回0，或错误码
 */
int ec_gen_keypair(EC_KEY* key)
{
    MPI_POINT P  = NULL;
    int       rc = 0;

    if (!key) return -EINVAL;

    P = mpi_point_new(0);
    if (!P) return -ENOMEM;

    rc = ec_gen_privkey(key);
    if (rc) goto free_p;

    mpi_ec_mul_point(P, key->priv, sm2_ctx.G, &sm2_ctx);

    rc = mpi_ec_get_affine(key->pub->x, key->pub->y, P, &sm2_ctx);
    if (rc) goto free_p;

    mpi_set_ui(key->pub->z, 1);

free_p:
    mpi_point_release(P);
    return rc;
}

/**
 * @brief 检查公钥有效性
 *
 * @param key 密钥
 * @return int 成功返回0，或错误码
 */
int ec_check_pubkey(EC_KEY* key)
{
    MPI_POINT point = NULL;
    MPI_POINT nP    = NULL;
    int       rc    = 0;

    if (!key) return -EINVAL;

    point = key->pub;

    nP = mpi_point_new(0);
    if (!nP) return -ENOMEM;

    /* a: 验证P不是无穷远点O */
    if (mpi_cmp_ui(point->z, 0) == 0) {
        rc = 1;
        goto free;
    }

    /* b: 验证 Yp∈[0, p-1]; */
    if (mpi_cmpabs(point->x, sm2_ctx.p) >= 0 ||
        mpi_cmpabs(point->y, sm2_ctx.p) >= 0) {
        rc = 2;
        goto free;
    }

    /* c: 验证y^2 == x^3 + ax + b (modp); 即点在曲线上 */
    if (!mpi_ec_curve_point(point, &sm2_ctx)) {
        rc = 4;
        goto free;
    }

    /* d:验证nP==O; */
    mpi_ec_mul_point(nP, sm2_ctx.n, point, &sm2_ctx);
    if (mpi_cmp_ui(nP->z, 0)) {
        rc = 5;
    }

free:
    mpi_point_release(nP);
    return rc;
}

/**
 * @brief 将密钥对转换为字节串
 *
 * @param key 密钥对
 * @param pubkey 公钥字节串[若存在]
 * @param privkey 私钥字节串[若存在]
 * @return int 成功返回0，或错误码
 */
int ec_key_to_buf(EC_KEY* key, unsigned char pubkey[SM2_PUBKEY_LEN],
                  unsigned char privkey[SM2_PRIVKEY_LEN])
{
    int          rc = 0;
    unsigned int nread;
    unsigned int buflen = SM2_PRIVKEY_LEN;
    if (privkey) mpi_read_buffer(key->priv, privkey, buflen, &nread, NULL);
    if (pubkey) {
        mpi_read_buffer(key->pub->x, pubkey, buflen, &nread, NULL);
        mpi_read_buffer(key->pub->y, pubkey + buflen, buflen, &nread, NULL);
    }
    return rc;
}